#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import sys, time, threading, csv, subprocess
import csv, os
from datetime import datetime
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

_PLOT_FIELDS = [
    ("gpu_util_percent", "GPU util (%)", "gpu_util"),
    ("gpu_mem_percent", "GPU mem (%)", "gpu_mem"),
    ("cpu_percent", "CPU util (%)", "cpu_percent"),
    ("ram_percent", "RAM util (%)", "ram_percent"),
]
try:
    import psutil
except Exception:
    psutil = None

# NVML if available
_nvml_loaded = False
try:
    import pynvml

    pynvml.nvmlInit()
    _nvml_loaded = True
except Exception:
    _nvml_loaded = False


def now_ts():
    return datetime.now().isoformat(timespec="seconds")


class ResourceMonitor:
    """Background sampler for GPU/CPU/RAM. Writes CSV periodically."""

    def __init__(self, csv_path, concurrency, need_header, interval=0.5, gpu_index=0):
        self.csv_path = csv_path
        self.interval = interval
        self.concurrency = concurrency
        self.need_header = need_header
        self.gpu_index = gpu_index
        self._stop = threading.Event()
        self._thr = None
        self.header = [
            "concurrency",
            "timestamp",
            "gpu_util_percent",
            "gpu_mem_percent",
            "cpu_percent",
            "ram_percent",
        ]

    def _gpu_metrics_nvml(self):
        try:
            h = pynvml.nvmlDeviceGetHandleByIndex(self.gpu_index)
            util = pynvml.nvmlDeviceGetUtilizationRates(h).gpu
            mem = pynvml.nvmlDeviceGetMemoryInfo(h)
            return util, int(mem.used / 1024 / 1024), int(mem.total / 1024 / 1024)
        except Exception:
            return None, None, None

    def _gpu_metrics_smi(self):
        try:
            out = subprocess.check_output(
                ["nvidia-smi",
                 "--query-gpu=utilization.gpu,memory.used,memory.total",
                 "--format=csv,noheader,nounits", "-i", str(self.gpu_index)],
                text=True
            ).strip()
            parts = out.split(",")
            return int(parts[0]), int(parts[1]), int(parts[2])
        except Exception:
            return None, None, None

    def _gpu_metrics(self):
        if _nvml_loaded:
            return self._gpu_metrics_nvml()
        return self._gpu_metrics_smi()

    def _cpu_ram(self):
        if psutil is None:
            return None, None, None
        cpu = psutil.cpu_percent(interval=None)
        vm = psutil.virtual_memory()
        return cpu, int(vm.used / 1024 / 1024), float(vm.percent)

    def start(self):
        if self.need_header:
            with open(self.csv_path, "w", encoding="utf-8", newline="") as f:
                w = csv.writer(f)
                w.writerow(self.header)
        # thread
        self._thr = threading.Thread(target=self._run, daemon=True)
        self._thr.start()

    def _run(self):
        while not self._stop.is_set():
            ts = now_ts()
            gu, gm_u, gm_t = self._gpu_metrics()  # util, used(MiB), total(MiB)
            cpu, ram_u, ram_p = self._cpu_ram()  # cpu%, used(MiB), ram%
            # 计算 GPU 显存占比（百分比）
            if gm_u is not None and gm_t:
                gm_percent = round(100.0 * gm_u / max(gm_t, 1), 3)
            else:
                gm_percent = None
            with open(self.csv_path, "a", encoding="utf-8", newline="") as f:
                import csv
                w = csv.writer(f)
                w.writerow([str(self.concurrency), ts, gu, gm_percent, cpu, ram_p])
            time.sleep(self.interval)

    def stop(self):
        self._stop.set()
        if self._thr:
            self._thr.join(timeout=2)


def _read_metrics_csv(path: str):
    xs, cols = [], {k: [] for k, _, _ in _PLOT_FIELDS}
    with open(path, "r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for r in reader:
            con = r.get("concurrency") or ""
            ts = r.get("timestamp") or ""
            try:
                xs.append("c" + con + "\n" + str(datetime.fromisoformat(ts)).split(" ")[-1])
            except Exception:
                xs.append(None)
            for key, _, _ in _PLOT_FIELDS:
                v = r.get(key)
                if v in (None, "", "None"):
                    cols[key].append(None)
                else:
                    try:
                        cols[key].append(float(v))
                    except:
                        cols[key].append(None)
    return xs, cols


def _moving_avg(seq, window: int):
    if not window or window <= 1: return seq
    out, buf = [], []
    for v in seq:
        if v is not None: buf.append(v)
        if len(buf) > window: buf.pop(0)
        out.append((sum(buf) / len(buf)) if buf else None)
    return out


def _plot_series(xs, ys, title: str, ylabel: str, out_path: str):
    fig = plt.figure(figsize=(10, 4.5))
    plt.plot(xs, ys)  # 不设颜色，单图
    plt.title(title)
    plt.xlabel("time")
    plt.ylabel(ylabel)
    plt.tight_layout()
    fig.savefig(out_path, dpi=140)
    plt.close(fig)


def _plot_multiple_series(xs, series, title: str, out_path: str):
    fig = plt.figure(figsize=(12, 6))
    for ys, label in series:
        plt.plot(xs, ys, label=label)
    plt.title(title)
    plt.xlabel("time")
    plt.ylabel("Metrics (%)")
    plt.legend()
    plt.tight_layout()
    fig.savefig(out_path, dpi=140)
    plt.close(fig)


def save_charts_from_metrics_csv(csv_path: str, outdir: str, prefix: str = "", ma: int = 1, title_suffix: str = ""):
    os.makedirs(outdir, exist_ok=True)
    xs, cols = _read_metrics_csv(csv_path)

    # 收集所有数据系列
    series = []
    for key, ylabel, fname in _PLOT_FIELDS:
        ys = _moving_avg(cols[key], ma)
        series.append((ys, ylabel))

    # 保存组合图表
    out_path = os.path.join(outdir, f"{prefix}all_metrics.png")
    _plot_multiple_series(xs, series, title=f"All Metrics over time{title_suffix}", out_path=out_path)
    print(f"[OK] Saved combined chart: {out_path}")
